import sys
import json
from llama_index.llms.ollama import Ollama
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.embeddings.fastembed import FastEmbedEmbedding
from llama_index.core import Settings
from llama_index.core.query_engine import NLSQLTableQueryEngine
import openai
import logging
from llama_index.llms.openai import OpenAI
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb
import uuid
from bs4 import BeautifulSoup
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.callbacks import CallbackManager
from llama_index.core.agent import AgentRunner
from llama_index.agent.openai import OpenAIAgent
from llama_index.agent.openai import OpenAIAgentWorker
from llama_index.embeddings.ollama import OllamaEmbedding
from llama_index.embeddings.openai import OpenAIEmbedding
import os
import re
import random
import html_to_json
from dotenv import load_dotenv
load_dotenv()
from llama_index.multi_modal_llms.ollama import OllamaMultiModal


llm = OllamaMultiModal(model="llava-llama3:latest", base_url="http://127.0.0.1:9090",request_timeout=600000.0)
from llama_index.core.program import MultiModalLLMCompletionProgram
from llama_index.core.output_parsers import PydanticOutputParser

from pathlib import Path
from llama_index.core import SimpleDirectoryReader
from PIL import Image
import matplotlib.pyplot as plt

# load as image documents
image_documents = SimpleDirectoryReader(r"E:\100rag\data\pic").load_data()

from pydantic import BaseModel


class Restaurant(BaseModel):
    """Data model for an restaurant."""

    restaurant: str
    food: str
    discount: str
    price: str
    rating: str
    review: str
    
prompt_template_str = """\
{query_str}

Return the answer as a Pydantic object. The Pydantic schema is given below:

"""
mm_program = MultiModalLLMCompletionProgram.from_defaults(
    output_parser=PydanticOutputParser(Restaurant),
    image_documents=image_documents,
    prompt_template_str=prompt_template_str,
    multi_modal_llm=llm,
    verbose=True,
)
response = mm_program(query_str="Can you summarize what is in the image?")
for res in response:
    print(res)